from pyscf import gto, scf, dft, mp, ci, cc
import numpy as np
import matplotlib.pyplot as plt

# Bohr radius
a0 = 0.52917721067

# 1 hartree energy = 27.211386245988eV
hartree = 27.211386245988

# 使用 STO-3G 构建 H2 分子
def H2(r:float):
    return gto.M(
        atom = f'H 0 0 0; H 0 0 {r}',
        basis = 'sto-3g'
    )

# 使用 STO-3G 构建 HeH+ 离子
def HeH(r:float):
    return gto.M(
        atom = f'He 0 0 0; H 0 0 {r}',
        basis = 'sto-3g',
        charge = 1
    )

def test_DFT(r, molecule):
    Erhf = []
    Elda = []
    Epbe = []
    Epbe0 = []
    for ri in r:
        mol = molecule(ri)
        Erhf.append(scf.RHF(mol).kernel())
        Elda.append(dft.RKS(mol, xc="lda").kernel())
        Epbe.append(dft.RKS(mol, xc="pbe").kernel())
        Epbe0.append(dft.RKS(mol, xc="pbe0").kernel())
    
    plt.figure()
    plt.plot(r / a0, Erhf, label="RHF")
    plt.plot(r / a0, Elda, label="LDA")
    plt.plot(r / a0, Epbe, label="PBE")
    plt.plot(r / a0, Epbe0, label="PBE0")
    plt.xlabel(r'$r/\mathrm{a.u.}$')
    plt.ylabel(r'$E_0/\mathrm{a.u.}$')
    plt.legend()
    plt.savefig(f"DFT-{molecule.__name__}.png")

def test_wave_function_methods(r, molecule):
    Erhf = []
    Euhf = []
    Eccsd = []
    Emp2 = []
    Ecisd = []
    for ri in r:
        print(f"ri = {ri}")
        mol = molecule(ri)
        Erhf.append(scf.RHF(mol).kernel())

        uhf_handle = scf.UHF(mol)
        dm_alpha, dm_beta = uhf_handle.get_init_guess()
        dm_beta[0,0] = 0  # 破坏对称性，参考：https://github.com/pyscf/pyscf/blob/master/examples/scf/32-break_spin_symm.py
        Euhf.append(uhf_handle.kernel((dm_alpha, dm_beta)))

        mymp = mp.MP2(uhf_handle).run()
        Emp2.append(mymp.e_tot)

        mycc = cc.CCSD(uhf_handle).run()
        Eccsd.append(mycc.e_tot)

        myci = ci.CISD(uhf_handle).run()
        Ecisd.append(myci.e_tot)
    
    r = r / a0
    plt.figure()
    plt.plot(r, Erhf, label="RHF")
    plt.plot(r, Euhf, label="UHF")
    plt.plot(r, Emp2, label="MP2")
    plt.plot(r, Ecisd, label="CISD")
    plt.plot(r, Eccsd, label="CCSD")
    plt.legend()
    plt.xlabel(r'$r/\mathrm{a.u.}$')
    plt.ylabel(r'$E_0/\mathrm{a.u.}$')
    plt.savefig(f'wave-function-{molecule.__name__}.png')

if __name__ == '__main__':
    # test_DFT(np.linspace(0.4, 2), H2)
    # test_DFT(np.linspace(0.5, 2.5), HeH)
    test_wave_function_methods(np.linspace(0.4, 2), H2)
    test_wave_function_methods(np.linspace(0.5, 2.5), HeH)